{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "d0d48696",
      "metadata": {
        "id": "d0d48696"
      },
      "source": [
        "## Install (Colab Only)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2d749094",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2d749094",
        "outputId": "3da4a3cd-def8-4b8b-8b01-19ef9e470309"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting pyepo\n",
            "  Downloading pyepo-0.4.0-py3-none-any.whl.metadata (1.2 kB)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from pyepo) (1.26.4)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from pyepo) (1.13.1)\n",
            "Collecting pathos (from pyepo)\n",
            "  Downloading pathos-0.3.3-py3-none-any.whl.metadata (11 kB)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from pyepo) (4.67.1)\n",
            "Requirement already satisfied: Pyomo>=6.1.2 in /usr/local/lib/python3.10/dist-packages (from pyepo) (6.8.2)\n",
            "Collecting gurobipy>=9.1.2 (from pyepo)\n",
            "  Downloading gurobipy-12.0.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (15 kB)\n",
            "Requirement already satisfied: scikit_learn in /usr/local/lib/python3.10/dist-packages (from pyepo) (1.6.0)\n",
            "Requirement already satisfied: torch>=1.13.1 in /usr/local/lib/python3.10/dist-packages (from pyepo) (2.5.1+cu121)\n",
            "Requirement already satisfied: ply in /usr/local/lib/python3.10/dist-packages (from Pyomo>=6.1.2->pyepo) (3.11)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo) (3.16.1)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo) (4.12.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo) (3.1.4)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo) (2024.10.0)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch>=1.13.1->pyepo) (1.3.0)\n",
            "Collecting ppft>=1.7.6.9 (from pathos->pyepo)\n",
            "  Downloading ppft-1.7.6.9-py3-none-any.whl.metadata (12 kB)\n",
            "Collecting dill>=0.3.9 (from pathos->pyepo)\n",
            "  Downloading dill-0.3.9-py3-none-any.whl.metadata (10 kB)\n",
            "Collecting pox>=0.3.5 (from pathos->pyepo)\n",
            "  Downloading pox-0.3.5-py3-none-any.whl.metadata (8.0 kB)\n",
            "Collecting multiprocess>=0.70.17 (from pathos->pyepo)\n",
            "  Downloading multiprocess-0.70.17-py310-none-any.whl.metadata (7.2 kB)\n",
            "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from scikit_learn->pyepo) (1.4.2)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.10/dist-packages (from scikit_learn->pyepo) (3.5.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.13.1->pyepo) (3.0.2)\n",
            "Downloading pyepo-0.4.0-py3-none-any.whl (45 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m45.1/45.1 kB\u001b[0m \u001b[31m1.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading gurobipy-12.0.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (14.4 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m14.4/14.4 MB\u001b[0m \u001b[31m21.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading pathos-0.3.3-py3-none-any.whl (82 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m82.1/82.1 kB\u001b[0m \u001b[31m4.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dill-0.3.9-py3-none-any.whl (119 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m119.4/119.4 kB\u001b[0m \u001b[31m6.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading multiprocess-0.70.17-py310-none-any.whl (134 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading pox-0.3.5-py3-none-any.whl (29 kB)\n",
            "Downloading ppft-1.7.6.9-py3-none-any.whl (56 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.8/56.8 kB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: gurobipy, ppft, pox, dill, multiprocess, pathos, pyepo\n",
            "Successfully installed dill-0.3.9 gurobipy-12.0.0 multiprocess-0.70.17 pathos-0.3.3 pox-0.3.5 ppft-1.7.6.9 pyepo-0.4.0\n"
          ]
        }
      ],
      "source": [
        "# install\n",
        "!pip install pyepo\n",
        "!pip install gurobipy"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Remove the problematic hook from Google Colab."
      ],
      "metadata": {
        "id": "GOVah1Ug1XGy"
      },
      "id": "GOVah1Ug1XGy"
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.meta_path = [hook for hook in sys.meta_path if not any(keyword in str(hook) for keyword in [\"google.colab\"])]"
      ],
      "metadata": {
        "id": "Znstbpcv1XO-"
      },
      "id": "Znstbpcv1XO-",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "ec3ca7c6",
      "metadata": {
        "id": "ec3ca7c6"
      },
      "source": [
        "# Optimization Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5a5b458f",
      "metadata": {
        "id": "5a5b458f"
      },
      "source": [
        "PyEPO ontains synthetic data generator and a dataset class ``optDataset`` to wrap data samples."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "984e785f",
      "metadata": {
        "id": "984e785f"
      },
      "source": [
        "## 1 Data Generator"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2854fb98",
      "metadata": {
        "id": "2854fb98"
      },
      "source": [
        "``pyepo.data`` includes synthetic datasets for three of the most classic optimization problems: the shortest path problem, the multi-dimensional knapsack problem, and the traveling salesperson problem.\n",
        "\n",
        "The synthetic datasets include features $\\mathbf{x}$ and cost coefficients $\\mathbf{c}$. The feature vector $\\mathbf{x}_i \\in \\mathbb{R}^p$ follows a standard multivariate Gaussian distribution $\\mathcal{N}(0, \\mathbf{I})$, and the corresponding cost $\\mathbf{c}_i \\in \\mathbb{R}^d$ comes from a polynomial function $f(\\mathbf{x}_i)$ multiplicated with a random noise $\\mathbf{\\epsilon}_i \\sim  U(1-\\bar{\\epsilon}, 1+\\bar{\\epsilon})$. The details of $f(\\mathbf{x}_i)$ can be seen [here](https://khalil-research.github.io/PyEPO/build/html/content/examples/data.html).\n",
        "\n",
        "In general, there are several parameters that users can control:\n",
        "- num_data ($n$): data size\n",
        "- num_features ($p$): feature dimension of costs $\\mathbf{c}$\n",
        "- deg ($deg$): polynomial degree of function $f(\\mathbf{x}_i)$\n",
        "- noise_width ($\\bar{\\epsilon}$):  noise half-width of $\\mathbf{\\epsilon}$\n",
        "- seed: random state seed to generate data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "546e4c22",
      "metadata": {
        "id": "546e4c22"
      },
      "source": [
        "### 1.1 Shortest Path"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3604f3e0",
      "metadata": {
        "id": "3604f3e0"
      },
      "source": [
        "$c_i^j = [\\frac{1}{{3.5}^{deg}} (\\frac{1}{\\sqrt{p}}(\\mathcal{B} \\mathbf{x}_i)_j + 3)^{deg} + 1] \\cdot \\epsilon_i^j$, where $\\mathcal{B}$ is a random matrix."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b013bd02",
      "metadata": {
        "id": "b013bd02"
      },
      "source": [
        "The following code is to generate data for the shortest path on the **5x5** grid network:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f49f3ae8",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f49f3ae8",
        "outputId": "d9a0e8eb-5d37-4ca1-9c61-b101f8606d7b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Auto-Sklearn cannot be imported.\n"
          ]
        }
      ],
      "source": [
        "import pyepo\n",
        "grid = (5,5) # grid size\n",
        "n = 1000 # number of data\n",
        "p = 5 # feature dimention\n",
        "deg = 4 # polynomial degree\n",
        "e = 0.5 # noise half-width"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a3c72c07",
      "metadata": {
        "id": "a3c72c07"
      },
      "outputs": [],
      "source": [
        "# generate data for grid network (features and costs)\n",
        "feats, costs = pyepo.data.shortestpath.genData(num_data=n+1000, num_features=p, grid=grid,\n",
        "                                               deg=deg, noise_width=e, seed=42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2a001df7",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2a001df7",
        "outputId": "53afe67c-c5a7-4aec-85cf-678e6413a374"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Features:\n",
            "[-0.68002472  0.2322537   0.29307247 -0.71435142  1.86577451]\n",
            "Costs:\n",
            "[0.35391761 1.07858994 1.01396442 0.66493631 0.29919599 0.19120229\n",
            " 0.92962621 0.24648609 0.33950669 0.19379419 0.78927149 0.52358102\n",
            " 0.66013164 1.0846318  0.73023344 0.53221019 0.34958448 0.33129203\n",
            " 1.4367641  0.72844442 1.39402493 0.89432676 1.03169003 0.46667478\n",
            " 0.60684515 1.7148205  1.62555298 2.13901473 0.375338   0.51937908\n",
            " 1.30751427 2.39109315 0.51398154 1.02980917 0.73931099 0.23779171\n",
            " 0.35521389 0.25666491 0.70956306 1.69988172]\n"
          ]
        }
      ],
      "source": [
        "# features\n",
        "print(\"Features:\")\n",
        "print(feats[0])\n",
        "# costs\n",
        "print(\"Costs:\")\n",
        "print(costs[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8cb8b5f9",
      "metadata": {
        "id": "8cb8b5f9"
      },
      "source": [
        "### 1.2 Knapsack"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f80ade8b",
      "metadata": {
        "id": "f80ade8b"
      },
      "source": [
        "Because we assume that the uncertain coefficients exist only on the objective function, the weights of items are fixed throughout the data. We define the number of items as $m$ and the dimension of resources is $k$."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "eaba994d",
      "metadata": {
        "id": "eaba994d"
      },
      "source": [
        "$c_i^j = \\lceil [\\frac{5}{{3.5}^{deg}} (\\frac{1}{\\sqrt{p}}(\\mathcal{B} \\mathbf{x}_i)_j + 3)^{deg} + 1] \\cdot \\epsilon_i^j \\rceil$, where $\\mathcal{B}$ is a random matrix."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "595beea0",
      "metadata": {
        "id": "595beea0"
      },
      "outputs": [],
      "source": [
        "import pyepo\n",
        "m = 32 # number of items\n",
        "k = 2 # resource dimension\n",
        "n = 1000 # number of data\n",
        "p = 5 # feature dimention\n",
        "deg = 4 # polynomial degree\n",
        "e = 0.5 # noise half-width"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "090b5a22",
      "metadata": {
        "id": "090b5a22"
      },
      "outputs": [],
      "source": [
        "# generate data for 2D kansack (weights, features and costs)\n",
        "weights, feats, costs = pyepo.data.knapsack.genData(num_data=n+1000, num_features=p, num_items=m,\n",
        "                                                    dim=k, deg=deg, noise_width=e, seed=42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "68b22b30",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "68b22b30",
        "outputId": "0b65fb1f-9d8a-4a85-8553-35e30af2196d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Weights of Items:\n",
            "[[4.02 7.35 6.48 5.7  4.06 3.71 4.88 3.2  4.02 4.21 7.66 5.14 6.3  7.58\n",
            "  3.87 6.72 3.99 6.59 4.51 4.3  4.49 6.08 5.57 6.43 7.91 7.13 5.93 6.85\n",
            "  4.91 7.43 5.76 4.6 ]\n",
            " [7.59 6.13 3.21 5.52 5.35 6.44 3.48 7.74 3.58 4.69 7.75 4.87 7.63 5.7\n",
            "  4.89 7.45 4.74 7.45 3.5  6.63 3.54 5.43 6.19 4.3  7.84 6.06 4.34 3.2\n",
            "  6.28 4.66 5.73 6.87]]\n",
            "Features:\n",
            "[-0.60139725 -0.84661436  0.81615272  0.47170729 -0.2854893 ]\n",
            "Values:\n",
            "[2. 5. 6. 4. 1. 2. 1. 4. 2. 2. 5. 5. 5. 5. 3. 2. 4. 2. 2. 2. 4. 2. 2. 4.\n",
            " 3. 2. 1. 3. 4. 2. 3. 5.]\n"
          ]
        }
      ],
      "source": [
        "# features\n",
        "print(\"Weights of Items:\")\n",
        "print(weights)\n",
        "print(\"Features:\")\n",
        "print(feats[0])\n",
        "# costs\n",
        "print(\"Values:\")\n",
        "print(costs[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "568a9cd4",
      "metadata": {
        "id": "568a9cd4"
      },
      "source": [
        "### 1.3 Traveling Salesperson"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dd739144",
      "metadata": {
        "id": "dd739144"
      },
      "source": [
        "For Traveling Salesperson (TSP),  number of nodes $m$ is the addtional parameter to generate data."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "657799ff",
      "metadata": {
        "id": "657799ff"
      },
      "source": [
        "The distance consists of two parts: one comes from Euclidean distance, the other derived from feature encoding. For Euclidean distance, we generate random coordinates. For feature encoding, it is $\\frac{1}{{3}^{deg - 1}} (\\frac{1}{\\sqrt{p}} (\\mathcal{B} \\mathbf{x}_i)_j + 3)^{deg} \\cdot \\epsilon_i^j$, where $\\mathcal{B}$ is a random matrix."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a9a4838d",
      "metadata": {
        "id": "a9a4838d"
      },
      "outputs": [],
      "source": [
        "import pyepo\n",
        "m = 20 # number of nodes\n",
        "n = 1000 # number of data\n",
        "p = 5 # feature dimention\n",
        "deg = 4 # polynomial degree\n",
        "e = 0.5 # noise half-width"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "33ef5f1e",
      "metadata": {
        "id": "33ef5f1e"
      },
      "outputs": [],
      "source": [
        "feats, costs = pyepo.data.tsp.genData(num_data=n+1000, num_features=p, num_nodes=m,\n",
        "                                      deg=deg, noise_width=e, seed=42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6ed1c3f5",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6ed1c3f5",
        "outputId": "98467d95-5b86-4266-bb0c-023e985ca986"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Features:\n",
            "[-1.14168911 -0.19365946 -0.71682232 -1.86653662 -0.08268069]\n",
            "Costs:\n",
            "[ 3.7469  3.5177  3.5663  3.2863 20.8808  3.5013  4.6786  3.663   9.4076\n",
            "  1.6293  8.884   2.8502  5.8908  4.8622 18.5146  5.7186  0.5103  3.1516\n",
            "  8.1697  6.4878 10.9322  0.766  12.8021  7.2463 16.3108  3.0362 11.5717\n",
            "  9.5445  3.901   4.9748  2.0201  2.4792  5.3137 26.4226 19.2363  2.3725\n",
            "  1.9566  6.9787  6.7527  6.1478  9.8816  4.8733  4.8745  2.1084  6.8356\n",
            "  7.9441  6.6348  3.2901 10.0526  4.4324  3.8492  6.6119  8.6355 22.0806\n",
            "  5.5779  1.3496  8.7577  5.0559  6.8631  6.8527  3.5512  3.4764  4.0971\n",
            "  4.0949  4.0859  7.7309  7.4587  1.2292 34.8099  4.9677 17.7628  3.1806\n",
            "  6.3381  3.5345 18.8379  6.9102  6.5611  3.818   4.2513 16.7784  4.5815\n",
            "  5.3674  3.3281 13.4974  2.304   4.8323  3.4517  7.0979  5.5551 10.3592\n",
            "  3.6593 15.9906  5.1407  4.9056  2.8864  5.8821  4.7064 16.2864  8.4832\n",
            "  6.8939  3.0092  6.5853  3.5005  6.8579  4.3257  2.9691  5.9071  4.8555\n",
            "  4.6766  4.2926  3.0147  4.5696  1.5417  2.2904  5.9068 13.253   7.2816\n",
            "  3.7608  2.6785  5.1225 12.4957  6.5661  3.4804  4.7496  5.1713  4.0051\n",
            " 10.6108  2.366   3.8686  2.858   2.1941  2.9987  5.616   4.2036  2.8274\n",
            "  1.7951  3.1929  1.9771  3.7794  7.1023  2.589   5.9645  3.5209  3.5754\n",
            " 23.3991  3.6     3.2637  4.7308  3.0393  3.0881 32.1364  6.1012  2.3523\n",
            " 10.6854  3.7003  1.3595  7.6433  5.9098  5.0089  4.3011 40.8034  4.4646\n",
            "  5.603   3.6141  6.6556  5.9318 13.4098  2.204  11.5992  1.972   5.4251\n",
            "  3.4477  4.9676  0.5672  4.2034  2.9343  4.3879  2.1213  1.4663  8.2608\n",
            "  1.6514 23.8687  5.0414  3.1683  3.1189 12.3226 19.7905  4.2913 35.3521\n",
            "  2.9655]\n"
          ]
        }
      ],
      "source": [
        "# features\n",
        "print(\"Features:\")\n",
        "print(feats[0])\n",
        "# costs\n",
        "print(\"Costs:\")\n",
        "print(costs[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1.4 Portfolio"
      ],
      "metadata": {
        "id": "gry3b3FNg_yF"
      },
      "id": "gry3b3FNg_yF"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let $\\bar{r}_{ij} = (\\frac{0.05}{\\sqrt{p}}(\\mathcal{B} \\mathbf{x}_i)_j + {0.1}^{\\frac{1}{deg}})^{deg}$.\n",
        "\n",
        "The expected return of the assets is defined as $\\bar{\\mathbf{r}}_i + \\mathbf{L} \\mathbf{f} + 0.01 \\tau \\mathbf{\\epsilon}$ and the covariance matrix is expressed $\\mathbf{L} \\mathbf{L}^{\\intercal} + (0.01 \\tau)^2 \\mathbf{I}$, where $\\mathcal{B}$, $\\mathbf{L}$, and $\\mathbf{f}$ and $\\mathbf{\\epsilon}$ are random variables."
      ],
      "metadata": {
        "id": "JtJ8HPnwhIGE"
      },
      "id": "JtJ8HPnwhIGE"
    },
    {
      "cell_type": "code",
      "source": [
        "import pyepo\n",
        "m = 50 # number of assets\n",
        "n = 1000 # number of data\n",
        "p = 4 # feature dimention\n",
        "deg = 4 # polynomial degree\n",
        "e = 1 # noise level"
      ],
      "metadata": {
        "id": "8zjp-Asziy8D"
      },
      "id": "8zjp-Asziy8D",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cov, feats, revs = pyepo.data.portfolio.genData(num_data=n+1000, num_features=p, num_assets=m,\n",
        "                                                deg=deg, noise_level=e, seed=42)"
      ],
      "metadata": {
        "id": "V9phXd3ti7Cz"
      },
      "id": "V9phXd3ti7Cz",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# covariance\n",
        "print(\"Covariance:\")\n",
        "print(cov)\n",
        "# features\n",
        "print(\"Features:\")\n",
        "print(feats[0])\n",
        "# Revenue\n",
        "print(\"Revenues:\")\n",
        "print(costs[0])"
      ],
      "metadata": {
        "id": "hm-LtIgVjKVE",
        "outputId": "4e5cb9e3-cfa6-47f0-92ea-3a9e46ead764",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "id": "hm-LtIgVjKVE",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Covariance:\n",
            "[[ 1.11661328e-04  1.04809132e-05  3.26874944e-06 ... -7.32011315e-06\n",
            "   2.31052654e-06  2.42490156e-06]\n",
            " [ 1.04809132e-05  1.10944311e-04  3.14384952e-06 ... -8.34225311e-06\n",
            "  -3.85284665e-07  3.65860588e-07]\n",
            " [ 3.26874944e-06  3.14384952e-06  1.09980349e-04 ... -5.52447038e-06\n",
            "   1.11220326e-06 -2.30590331e-06]\n",
            " ...\n",
            " [-7.32011315e-06 -8.34225311e-06 -5.52447038e-06 ...  1.10129657e-04\n",
            "  -1.57470787e-06  4.22118562e-06]\n",
            " [ 2.31052654e-06 -3.85284665e-07  1.11220326e-06 ... -1.57470787e-06\n",
            "   1.07591972e-04  1.05277517e-07]\n",
            " [ 2.42490156e-06  3.65860588e-07 -2.30590331e-06 ...  4.22118562e-06\n",
            "   1.05277517e-07  1.06530748e-04]]\n",
            "Features:\n",
            "[ 1.30547881  0.02100384  0.68195297 -0.31026676]\n",
            "Revenues:\n",
            "[ 3.7469  3.5177  3.5663  3.2863 20.8808  3.5013  4.6786  3.663   9.4076\n",
            "  1.6293  8.884   2.8502  5.8908  4.8622 18.5146  5.7186  0.5103  3.1516\n",
            "  8.1697  6.4878 10.9322  0.766  12.8021  7.2463 16.3108  3.0362 11.5717\n",
            "  9.5445  3.901   4.9748  2.0201  2.4792  5.3137 26.4226 19.2363  2.3725\n",
            "  1.9566  6.9787  6.7527  6.1478  9.8816  4.8733  4.8745  2.1084  6.8356\n",
            "  7.9441  6.6348  3.2901 10.0526  4.4324  3.8492  6.6119  8.6355 22.0806\n",
            "  5.5779  1.3496  8.7577  5.0559  6.8631  6.8527  3.5512  3.4764  4.0971\n",
            "  4.0949  4.0859  7.7309  7.4587  1.2292 34.8099  4.9677 17.7628  3.1806\n",
            "  6.3381  3.5345 18.8379  6.9102  6.5611  3.818   4.2513 16.7784  4.5815\n",
            "  5.3674  3.3281 13.4974  2.304   4.8323  3.4517  7.0979  5.5551 10.3592\n",
            "  3.6593 15.9906  5.1407  4.9056  2.8864  5.8821  4.7064 16.2864  8.4832\n",
            "  6.8939  3.0092  6.5853  3.5005  6.8579  4.3257  2.9691  5.9071  4.8555\n",
            "  4.6766  4.2926  3.0147  4.5696  1.5417  2.2904  5.9068 13.253   7.2816\n",
            "  3.7608  2.6785  5.1225 12.4957  6.5661  3.4804  4.7496  5.1713  4.0051\n",
            " 10.6108  2.366   3.8686  2.858   2.1941  2.9987  5.616   4.2036  2.8274\n",
            "  1.7951  3.1929  1.9771  3.7794  7.1023  2.589   5.9645  3.5209  3.5754\n",
            " 23.3991  3.6     3.2637  4.7308  3.0393  3.0881 32.1364  6.1012  2.3523\n",
            " 10.6854  3.7003  1.3595  7.6433  5.9098  5.0089  4.3011 40.8034  4.4646\n",
            "  5.603   3.6141  6.6556  5.9318 13.4098  2.204  11.5992  1.972   5.4251\n",
            "  3.4477  4.9676  0.5672  4.2034  2.9343  4.3879  2.1213  1.4663  8.2608\n",
            "  1.6514 23.8687  5.0414  3.1683  3.1189 12.3226 19.7905  4.2913 35.3521\n",
            "  2.9655]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "35d5ee2d",
      "metadata": {
        "id": "35d5ee2d"
      },
      "source": [
        "## 2 Introduction to optDataset"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "82c49072",
      "metadata": {
        "id": "82c49072"
      },
      "source": [
        "``optDataset`` is PyTorch Dataset, which stores the features and their corresponding costs of the objective function, and **solves optimization problems to get optimal solutions and optimal objective values**.\n",
        "\n",
        "``optDataset`` is **not** necessary for training with PyEPO, but it can be easier to obtain optimal solutions and objective values when they are not available in the original data.\n",
        "\n",
        "As the following example, ``optDataset`` and Pytorch ``DataLoader`` wrap the data samples, which can make the model training cleaner and more organized."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "98c8330c",
      "metadata": {
        "id": "98c8330c"
      },
      "source": [
        "### 2.1 Generate Data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1b07b942",
      "metadata": {
        "id": "1b07b942"
      },
      "source": [
        "We generate data for the shortest path on the 5x5 grid network first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8c8abdcb",
      "metadata": {
        "id": "8c8abdcb"
      },
      "outputs": [],
      "source": [
        "# grid size\n",
        "grid = (5,5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "448296e5",
      "metadata": {
        "id": "448296e5"
      },
      "outputs": [],
      "source": [
        "# generate data\n",
        "num_data = 1000 # number of data\n",
        "num_feat = 5 # size of feature\n",
        "deg = 4 # polynomial degree\n",
        "e = 0.5 # noise width\n",
        "feats, costs = pyepo.data.shortestpath.genData(num_data+1000, num_feat, grid, deg, e, seed=42)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "25b0ad04",
      "metadata": {
        "id": "25b0ad04"
      },
      "source": [
        "### 2.2 Build OptModel"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ac660eda",
      "metadata": {
        "id": "ac660eda"
      },
      "source": [
        "> \"PyTorch provides two data primitives: ``Dataset`` and ``DataLoader`` that allow you to use pre-loaded datasets as well as your own data. Dataset stores the samples and their corresponding labels, and DataLoader wraps an iterable around the Dataset to enable easy access to the samples. \"  -- PyTorch Documentation\n",
        "\n",
        "``optDataset`` is extended from PyTorch ``Dataset``. In order to obtain optimal solutions, ``optDataset`` requires the corresponding ``optModel`` is a module of PyEPO library, which is designed as a container for any \"black box\" solver. The tutorial on ``optModel`` is [here](https://github.com/khalil-research/PyEPO/blob/main/notebooks/01%20Optimization%20Model.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "da1c5dc3",
      "metadata": {
        "id": "da1c5dc3"
      },
      "source": [
        "Here we load the pre-defined [``pyepo.model.grb.shortestPathModel``](https://khalil-research.github.io/PyEPO/build/html/content/examples/model.html#pre-defined-models), which uses Gurobi to build a linear program."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "19fd4e0d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "19fd4e0d",
        "outputId": "53de1cf7-c4f8-4ec8-8802-9e4a4558c312"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Restricted license - for non-production use only - expires 2026-11-23\n"
          ]
        }
      ],
      "source": [
        "from pyepo.model.grb import shortestPathModel\n",
        "# init model\n",
        "optmodel = shortestPathModel(grid)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0d3c9171",
      "metadata": {
        "id": "0d3c9171"
      },
      "source": [
        "### 2.3 Train-Test Split"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0fb0fb74",
      "metadata": {
        "id": "0fb0fb74"
      },
      "source": [
        "Split features and costs into random train and test subsets, where the test size is 1000."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c0c1b954",
      "metadata": {
        "id": "c0c1b954"
      },
      "outputs": [],
      "source": [
        "# split train test data\n",
        "from sklearn.model_selection import train_test_split\n",
        "x_train, x_test, c_train, c_test = train_test_split(feats, costs, test_size=1000, random_state=42)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "80edb4ce",
      "metadata": {
        "id": "80edb4ce"
      },
      "source": [
        "### 2.4 Wrap Features and Costs with optDataset"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f24c1136",
      "metadata": {
        "id": "f24c1136"
      },
      "source": [
        "``optDataset`` accepts features features $\\mathbf{x}$ and cost coefficients $\\mathbf{c}$, then ultilizes ``optModel`` to find optimal solutions $\\mathbf{w}^*$ and objective values $\\mathbf{z}^*$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "56ce885b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "56ce885b",
        "outputId": "1019f581-c277-4d75-ca1f-5a8757a5f608"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizing for optDataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1000/1000 [00:04<00:00, 227.21it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizing for optDataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1000/1000 [00:03<00:00, 316.80it/s]\n"
          ]
        }
      ],
      "source": [
        "# get optDataset\n",
        "dataset_train = pyepo.data.dataset.optDataset(optmodel, x_train, c_train)\n",
        "dataset_test = pyepo.data.dataset.optDataset(optmodel, x_test, c_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "de8a3de5",
      "metadata": {
        "id": "de8a3de5"
      },
      "source": [
        "It takes time for solving optimization problem per instance."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e84ce621",
      "metadata": {
        "id": "e84ce621"
      },
      "source": [
        "The ``optDataset`` contains 4 attributes: **features**, **costs**, (optimal) **solutions**, and (optimal) **objective values**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f98646ae",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f98646ae",
        "outputId": "888c0ee2-0f8b-4ccc-acff-dd88f941ec38"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[ 0.68356932, -1.36595571,  1.21194399,  0.26125053, -0.36927714],\n",
              "       [-2.8321556 , -0.45115886,  0.5517408 ,  1.20026175, -0.46316136],\n",
              "       [-0.85238677,  0.47536561,  0.63245422, -0.47417818, -0.77177196],\n",
              "       [ 0.72894119, -0.69235077, -0.65903927, -0.57410068,  0.5364139 ],\n",
              "       [-1.23252331,  0.55229994,  0.62563093, -0.69677182,  0.58202657]])"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ],
      "source": [
        "# features\n",
        "dataset_train.feats[:5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d3f978e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9d3f978e",
        "outputId": "b747a7b9-e720-4c70-8c6e-9eddf98ab40b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.66423204, 1.31570609, 0.39314494, 1.56369511, 0.75385019,\n",
              "        1.31562851, 0.64220475, 0.93906923, 0.91311846, 2.04161427,\n",
              "        0.74907072, 0.5318552 , 0.68541336, 1.23490138, 0.93002909,\n",
              "        0.44799287, 0.25995308, 0.34913271, 0.41296698, 0.31735299,\n",
              "        0.21611971, 0.86564175, 0.8231998 , 0.21084682, 0.34851563,\n",
              "        0.5123044 , 0.64182268, 0.39734889, 0.5353895 , 0.58975985,\n",
              "        0.82311134, 0.22680251, 0.78549927, 0.87061416, 1.1983596 ,\n",
              "        0.2857453 , 1.45720774, 0.99978659, 0.30051617, 0.29062359],\n",
              "       [1.36202037, 1.60539413, 0.43680829, 0.43538718, 0.09457201,\n",
              "        0.20794562, 0.15560289, 0.10297456, 1.32280424, 0.17186773,\n",
              "        0.16530315, 0.06899043, 1.08747827, 0.03893738, 0.1008939 ,\n",
              "        0.02468207, 0.05977335, 1.20703588, 0.90134767, 0.26220181,\n",
              "        0.6374565 , 0.6714728 , 0.8947684 , 0.11022986, 0.02432527,\n",
              "        0.50323957, 0.54521407, 0.61761709, 0.08326333, 0.45347036,\n",
              "        0.04713232, 0.28148922, 0.68087899, 0.04876477, 0.09798262,\n",
              "        0.47139562, 0.83157447, 0.21111578, 0.79249161, 0.96827449],\n",
              "       [0.56509641, 0.1976001 , 0.88585206, 1.08069606, 0.32171541,\n",
              "        0.36611625, 0.15136523, 0.24740828, 0.22692131, 0.37700542,\n",
              "        0.3346859 , 0.3654156 , 0.95889515, 0.29711513, 0.1778639 ,\n",
              "        0.41939872, 0.44285341, 1.0000998 , 0.67106361, 0.82278396,\n",
              "        0.21169908, 0.64945546, 0.19477406, 0.14709023, 0.46549561,\n",
              "        0.86536996, 0.36663587, 0.35795097, 0.24298385, 1.07962878,\n",
              "        0.22336928, 0.90596097, 1.14643841, 0.25712994, 0.28179436,\n",
              "        0.68184758, 0.44373398, 0.2703023 , 0.43358737, 0.37312498],\n",
              "       [0.10794211, 0.33086822, 0.12473134, 0.4310949 , 1.02743299,\n",
              "        0.33948593, 0.77776785, 0.65206391, 0.22029336, 0.48640155,\n",
              "        0.37302538, 0.92430436, 0.44496363, 0.90479863, 0.47967375,\n",
              "        0.71608741, 0.61975307, 0.17753697, 0.16756215, 0.25667334,\n",
              "        0.46031177, 0.32373011, 0.38150178, 0.57426177, 0.60586239,\n",
              "        0.18689977, 0.59457171, 0.18590674, 0.80583248, 0.31708958,\n",
              "        0.96396724, 0.42087245, 0.11504448, 1.47034502, 0.70771891,\n",
              "        0.20560216, 0.29480205, 0.25587608, 0.12236201, 0.12280357],\n",
              "       [0.76456028, 0.75042964, 0.85989263, 0.58400794, 0.24933744,\n",
              "        0.25228797, 0.16888079, 0.2072617 , 0.37507889, 0.24028514,\n",
              "        0.46217285, 0.13744926, 1.07518581, 0.43993511, 0.25914406,\n",
              "        0.39713587, 0.38890696, 0.73725914, 0.6670779 , 0.5834132 ,\n",
              "        0.71862985, 1.00280771, 0.68093803, 0.23410139, 0.21084228,\n",
              "        1.60189634, 0.7013606 , 0.86136704, 0.28926341, 0.73420527,\n",
              "        0.36889941, 1.11041905, 1.02906701, 0.26626439, 0.49535779,\n",
              "        0.73181932, 0.4289328 , 0.39177901, 0.58704525, 0.59780448]])"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ],
      "source": [
        "# costs\n",
        "dataset_train.costs[:5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e3b07592",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e3b07592",
        "outputId": "0952b4c1-a3ed-4292-e613-436087c3223d"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,\n",
              "        0., 0., 1., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.,\n",
              "        0., 1., 0., 0., 0., 0., 1., 1.],\n",
              "       [0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 1.,\n",
              "        0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 1., 0.,\n",
              "        0., 0., 0., 1., 0., 0., 0., 0.],\n",
              "       [0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0.,\n",
              "        0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
              "        0., 1., 0., 0., 0., 0., 1., 1.],\n",
              "       [1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.,\n",
              "        0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
              "        1., 0., 0., 0., 0., 1., 1., 1.],\n",
              "       [0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0.,\n",
              "        0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
              "        0., 1., 0., 0., 0., 0., 1., 1.]])"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ],
      "source": [
        "# solutions\n",
        "dataset_train.sols[:5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1231e1ca",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1231e1ca",
        "outputId": "5d3e6e37-9d1b-45aa-81b2-a389be27184d"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[4.52934109],\n",
              "       [1.45274853],\n",
              "       [2.33050111],\n",
              "       [2.11744969],\n",
              "       [2.72324556]])"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "# objective values\n",
        "dataset_train.objs[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2f703e42",
      "metadata": {
        "id": "2f703e42"
      },
      "source": [
        "### 2.5 Set PyTorch DataLoader"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8912b782",
      "metadata": {
        "id": "8912b782"
      },
      "source": [
        "PyTorch ``DataLoader`` combines a dataset and a sampler, and provides an iterable over the given dataset. We should set the **batch size**, which is a number of samples processed before the model is updated."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "23fb2e94",
      "metadata": {
        "id": "23fb2e94"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "batch_size = 32\n",
        "loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)\n",
        "loader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d4bf5b77",
      "metadata": {
        "id": "d4bf5b77"
      },
      "source": [
        "To iterate ``DataLoader``, we can obtain a batch of **features**, **costs**, (optimal) **solutions**, and (optimal) **objective values**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ce529f39",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ce529f39",
        "outputId": "6559468b-3934-4ded-e343-0ee55b3659a4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([32, 5])\n",
            "torch.Size([32, 40])\n",
            "torch.Size([32, 40])\n",
            "torch.Size([32, 1])\n"
          ]
        }
      ],
      "source": [
        "for x, c, w, z in loader_train:\n",
        "    # shape of features batch\n",
        "    print(x.shape)\n",
        "    # shape of true costs batch\n",
        "    print(c.shape)\n",
        "    # shape of true optimal solutions batch\n",
        "    print(w.shape)\n",
        "    # shape of true optimal objective values batch\n",
        "    print(z.shape)\n",
        "    break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4dc28398",
      "metadata": {
        "id": "4dc28398"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}